#!/home/up/anaconda/envs/LPR1/bin/python
# encoding: utf-8
import rospy
import sys
import time
from std_msgs.msg import String
from geometry_msgs.msg import Point, PoseStamped, PointStamped
from sensor_msgs.msg import Image, CameraInfo
import numpy as np

sys.path.remove('/opt/ros/melodic/lib/python2.7/dist-packages')
import cv2
from cv_bridge import CvBridge, CvBridgeError


class image_transfer:
    def __init__(self):
        # param
        self.rate = rospy.Rate(100)
        self.desired_border_x = 30
        self.desired_border_y = 30


        self.count = 0
        self.count_color = 0
        self.count_yolo = 0
        self.image = None
        self.color_image = None
        self.depth_image = None
        self.TimeStamp = None
        self.p_in_pixel = [320, 240]    # 需要得到真实三维信息的像素点（x, y)
        self.border = [200, 400, 100, 300]  # [x_min,x_max,y_min,y_max]
        self.border_shrink = self.border
        self.p_mean = [0, 0, 0]         # 质心的坐标[x_mean,y_mean,z_mean]

        '''
        ros initial
        '''
        # rospy.init_node("image_transfer", anonymous=True)

        '''
        ros subscriber
        '''
        self.ros_yolo_corner_sub = rospy.Subscriber("/yolo_target_corner", PoseStamped, self.ros_yolo_corner_callback)
        self.ros_img_sub = rospy.Subscriber("/d435/aligned_depth_to_color/image_raw", Image, self.ros_img_callback)
        self.ros_img_color_sub = rospy.Subscriber("/d435/color/image_raw", Image, self.ros_img_color_callback)
        '''
        ros Publisher
        '''
        self.point_3d_pub = rospy.Publisher('/targetincamera_3d', PointStamped, queue_size=1)

    def ros_yolo_corner_callback(self, msg):
        if abs(msg.pose.orientation.x - msg.pose.orientation.z) != 0:
            self.border[0] = min(int(msg.pose.orientation.x),int(msg.pose.orientation.z))  # x min
            self.border[1] = max(int(msg.pose.orientation.x),int(msg.pose.orientation.z))  # x max
            self.border[2] = min(int(msg.pose.orientation.y),int(msg.pose.orientation.w))  # y min
            self.border[3] = max(int(msg.pose.orientation.y),int(msg.pose.orientation.w))  # y max
            self.count_yolo = 1
            # self.cal_least_square()
            self.TimeStamp = msg.header
            print('get yolo result', self.border)
        else:
            print('result is zero')

    def ros_img_callback(self, image_msg):
        bridge = CvBridge()
        # self.image = bridge.imgmsg_to_cv2(image_msg, "bgr8")
        self.depth_image = bridge.imgmsg_to_cv2(image_msg, '16UC1')
        self.count = 1

    def ros_img_color_callback(self, image_msg):
        bridge = CvBridge()
        self.color_image = bridge.imgmsg_to_cv2(image_msg, "bgr8")
        self.count_color = 1

    def cal_least_square(self):
        t1 = time.time()
        depth_image = self.depth_image

        # shrink the calculater area
        self.shrink_border()

        # get the border
        x_min = int(self.border_shrink[0])
        x_max = int(self.border_shrink[1])
        y_min = int(self.border_shrink[2])
        y_max = int(self.border_shrink[3])
        # print('depth image', np.size(depth_image,0))
        # print('depth image 2', np.size(depth_image,1))
        # print('self.border',self.border)

        # -------------------------for z in camera coordinate ---------------
        # get the depth information from image
        depth_mat = depth_image[y_min:y_max + 1, x_min:x_max + 1]  # need to transpose
        depth_array = depth_mat.reshape((1, -1))  # make it a row vector

        # filter the depth array
        idz = depth_array != 0  # find the zero point id
        depth_filter = depth_array[idz]  # # delete 0
        # print('z_array_filter', depth_filter)

        # if there is np useful point in the depth information
        if np.size(depth_filter) == 0:  # if there are no cloud points
            print('!!!there is no useful cloud point!!!')
            return

        # calculate mean depth
        z_mean = np.mean(depth_filter)
        t2 = time.time()
        # print('step1', t2-t1)

        # -------------------------for x and y in camera coordinate ---------------
        # initial the x and y vector
        x_line = np.arange(x_min, x_max+1).reshape((-1, 1))
        y_line = np.arange(y_min, y_max+1).reshape((1, -1))
        x_length = np.size(x_line)
        y_length = np.size(y_line)

        x_array = np.dot(x_line, np.ones((1, y_length)))
        y_array = np.dot(np.ones((x_length, 1)), y_line)
        x_array = x_array.reshape((1, -1))
        y_array = y_array.reshape((1, -1))
        t_g = time.time()
        # print('step2', t_g-t2)

        # calculate x and y in camera coordinate
        z_array = np.array(depth_array)
        # print('border',[x_min,x_max,y_min,y_max])
        # print('x_size', np.size(np.array(x_array)))
        # print('z_size', np.size(z_array))

        x_cal = (np.array(x_array) - ppx) / fx * z_array
        y_cal = (np.array(y_array) - ppy) / fy * z_array
        t3 = time.time()
        # print('step3', t3-t_g)

        # 去0
        id_x = x_cal != 0
        x_cal_filter = x_cal[id_x]
        id_y = y_cal != 0
        y_cal_filter = y_cal[id_y]
        t4 = time.time()
        # print('step4', t4 - t3)

        x_mean = np.mean(x_cal_filter)
        y_mean = np.mean(y_cal_filter)
        self.p_mean = [x_mean, y_mean, z_mean]
        # print('---------------------------------------')
        # print('mean', self.p_mean)

        t_end = time.time()
        # print('---------------------------------------')
        # print('spend time', t_end - t1)
        return x_mean, y_mean, z_mean

    def draw_in_pixel(self, img_color):
        x_min = self.border_shrink[0]
        x_max = self.border_shrink[1]
        y_min = self.border_shrink[2]
        y_max = self.border_shrink[3]
        cv2.rectangle(img_color, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
        # cv2.circle(img_color, (p_in_pixel[0], p_in_pixel[1]), 8, [255, 0, 255], thickness=-1)
        # cv2.putText(img_color, "Dis:" + str(dis) + " m", (40, 40), cv2.FONT_HERSHEY_SIMPLEX, 1.2, [0, 0, 255])
        # cv2.putText(img_color, "PIXEL_LEFT_X:" + str(self.border[0]) + " PIXEL_LEFT_Y:"+str(self.border[2]),
        #             (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1.2, [255, 0, 0])
        # cv2.putText(img_color, "PIXEL_RIGHT_X:" + str(self.border[1]) + "PIXEL_RIGHT_Y:"+str(self.border[3]),
        #             (40, 40), cv2.FONT_HERSHEY_SIMPLEX, 1.2, [255, 0, 0])
        cv2.putText(img_color, "X:" + str(self.p_mean[0]) + " m", (80, 80), cv2.FONT_HERSHEY_SIMPLEX, 1.2,
                    [255, 255, 255])
        cv2.putText(img_color, "Y:" + str(self.p_mean[1]) + " m", (80, 120), cv2.FONT_HERSHEY_SIMPLEX, 1.2,
                    [255, 255, 255])
        cv2.putText(img_color, "Z:" + str(self.p_mean[2]) + " m", (80, 160), cv2.FONT_HERSHEY_SIMPLEX, 1.2,
                    [255, 255, 255])
        return img_color

    def shrink_border(self):
        x_min = self.border[0]
        x_max = self.border[1]
        y_min = self.border[2]
        y_max = self.border[3]
        dx = x_max - x_min
        dy = y_max - y_min
        self.border_shrink = self.border
        if dx > self.desired_border_x:
            self.border_shrink[0] = int((x_max+x_min)/2 - self.desired_border_x/2)
            self.border_shrink[1] = int((x_max+x_min)/2 + self.desired_border_x/2)
        if dy > self.desired_border_y:
            self.border_shrink[2] = int((y_max+y_min)/2 - self.desired_border_y/2)
            self.border_shrink[3] = int((y_max+y_min)/2 + self.desired_border_y/2)

    def filter_depth(self):
        ""

    def task(self):
        while not rospy.is_shutdown():

            if self.count == 0 or self.count_color == 0:# doubt count_color need to be reset as 1 after process
                # print('count', self.count)
                continue

            self.cal_least_square()
            depth_image = self.draw_in_pixel(self.color_image)
            cv2.imshow('1', depth_image)
            # cv2.imshow('1', self.color_image)
            cv2.waitKey(3)

            # publish
            point_in_cam = PointStamped()
            point_in_cam.header = self.TimeStamp
            point_in_cam.point.x = self.p_mean[0]
            point_in_cam.point.y = self.p_mean[1]
            point_in_cam.point.z = self.p_mean[2]
            self.point_3d_pub.publish(point_in_cam)



if __name__ == '__main__':
    global fx, fy, ppx, ppy

    fx = 1367.017578125
    fy = 1368.4324951171875
    ppx = 978.2013549804688
    ppy = 523.9000854492188

    rospy.init_node("image_transfer", anonymous=True)
    image_demo1 = image_transfer()
    image_demo1.task()



